import os
import pandas as pd
import numpy as np
import pickle

class Factor:
    """因子基类"""
    def __init__(self, name):
        self.name = name
    
    def compute(self, data):
        """计算因子方法"""
        raise NotImplementedError
    
class ret(Factor):
    """区间收益率因子"""
    def __init__(self, n_days, name):
        super().__init__(f"{name}_{n_days}d")
        self.n_days = n_days
    
    def compute(self, dfs):
        """
        计算逻辑：
        1. 按股票代码分组
        2. 计算n日收益率
        3. 返回因子值
        """
        # 加载初始因子
        data_close=dfs['close']

        # 计算n日收益率
        data_close_pre = data_close.shift(self.n_days)
        factor=(data_close / data_close_pre) - 1
        return factor.replace({-np.inf:np.nan,np.inf:np.nan})

class exp_wgt_return_3m:
    def __init__(self, close_df, turn_df, lambda_value=0.05):
        """
        :param lambda_value: 指数衰减因子，默认为 0.05
        """
        self.close_df = close_df
        self.turn_df = turn_df
        self.lambda_value = lambda_value

    def calculate_daily_return(self):
        """
        计算每日收益率
        """
        return self.close_df.pct_change()

    def calculate_63day_return_weighted_factor_single_stock(self, stock_close, stock_turn):
        """
        计算单个股票的近 63 日收益率以换手率指数衰减加权的因子值
        """
        stock_return = stock_close.pct_change()
        factor_values = []
        for i in range(len(stock_return)):
            if i < 63:
                factor_values.append(np.nan)
            else:
                recent_63_days_return = stock_return.iloc[i - 63:i]
                recent_63_days_turn = stock_turn.iloc[i - 63:i]
                days = np.arange(63)
                weights = np.exp(-self.lambda_value * (63 - days - 1))
                numerator = (recent_63_days_return * recent_63_days_turn * weights).sum()
                denominator = (recent_63_days_turn * weights).sum()
                if denominator == 0:
                    factor_values.append(np.nan)
                else:
                    factor_value = numerator / denominator
                    factor_values.append(factor_value)
        return pd.Series(factor_values, index=stock_return.index)

    def calculate_63day_return_weighted_factor_all_stocks(self):
        """
        计算所有股票的近 63 日收益率以换手率指数衰减加权的因子值
        """
        factor_df = pd.DataFrame(index=self.close_df.index)
        for stock in self.close_df.columns:
            stock_close = self.close_df[stock]
            stock_turn = self.turn_df[stock]
            factor_df[stock] = self.calculate_63day_return_weighted_factor_single_stock(stock_close, stock_turn)
        return factor_df


#收益率近21日标准差

class std(Factor):

    def __init__(self, close_df, turn_df, lambda_value=0.05):
        self.close_df = close_df
        self.turn_df = turn_df
        self.lambda_value = lambda_value

    def calculate_daily_return(self):
        """
        计算每日收益率
        """
        return self.close_df.pct_change()

    def calculate_21day_return_std_all_stocks(self):
        """
        计算所有股票的近 21 日收益率标准差
        """
        daily_returns = self.calculate_daily_return()
        std_df = daily_returns.rolling(window=21).std()
        return std_df

# 计算相关系数
def calculate_correlation(x, y):
    return x.corr(y)



# 定义衰减因子，这里以 0.85 为例,目的是聚焦短期影响
alpha = 0.85
# 计算指数衰减加权收益率
def calculate_weighted_return(return_window, turn_window):
    window_length = len(return_window)
    # 计算指数衰减权重
    weights = np.array([alpha ** i for i in range(window_length)][::-1])
    # 归一化权重
    weights = weights / weights.sum()
    # 计算加权收益率
    weighted_return = (return_window * turn_window * weights).sum()
    return weighted_return


    



    

